import math
import torch
from math import exp
import torch.nn.functional as F
from torch.utils.data.dataset import Dataset

class Image_Quality_Metric():
	def __init__(self):
		print('Metric Initialized...')

	def gaussian(self, window_size, sigma):
	    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
	    return gauss/gauss.sum()

	def create_window(self, window_size, channel):
	    _1D_window = self.gaussian(window_size, 1.5).unsqueeze(1)
	    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
	    window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
	    return window

	def _ssim(self, img1, img2, window, window_size, channel, size_average = True):
	    mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel)
	    mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel)

	    mu1_sq = mu1.pow(2)
	    mu2_sq = mu2.pow(2)
	    mu1_mu2 = mu1*mu2

	    sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq
	    sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq
	    sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2

	    C1 = 0.01**2
	    C2 = 0.03**2

	    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))

	    if size_average:
	        return ssim_map.mean()
	    else:
	        return ssim_map.mean(1).mean(1).mean(1)

	def calssim(self, img1, img2, window_size = 11, size_average = True):
	    (_, channel, _, _) = img1.size()
	    window = self.create_window(window_size, channel)
	    
	    if img1.is_cuda:
	        window = window.cuda(img1.get_device())
	    window = window.type_as(img1)
	    
	    return self._ssim(img1, img2, window, window_size, channel, size_average)

	def rgb2y(self, img, remove_border):
		img = img[..., remove_border:-remove_border, remove_border:-remove_border]
		if img.size(1) > 1:
			gray_coeffs = [65.738, 129.057, 25.064]
			convert = img.new_tensor(gray_coeffs).view(1, 3, 1, 1) / 256
			img = 16.0+img.mul(convert).sum(dim=1)
			img = torch.round(img)
		return img/255.0

	def ssim(self, imgS, imgG, remove_border):
		imgS = self.rgb2y(imgS, remove_border)
		imgG = self.rgb2y(imgG, remove_border)
		return self.calssim(imgS, imgG)

	def psnr(self, imgS, imgG, remove_border):
		imgS = self.rgb2y(imgS, remove_border)
		imgG = self.rgb2y(imgG, remove_border)
		diff = imgS - imgG
		mse = diff.pow(2).mean()
		return -10 * math.log10(mse)

